package com.john.javacv.ffmpeg;

import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.FrameGrabber;
import org.bytedeco.javacv.Java2DFrameConverter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.awt.image.BufferedImage;

public class VideoCaptureUtil {
    private static final Logger LOGGER = LoggerFactory.getLogger(VideoCaptureUtil.class);

    /**
     * 截取视频图片
     *
     * @param fileUrl    视频文件地址 本地文件地址，或者是网络地址
     * @param frameIndex 要获取的帧序号 >=0的值 如果指定的序号位置没有图片，将使用后面最近的一个。
     * @return 帧图片对象
     */
    public static VideoCaptureImageResult capture(String fileUrl, int frameIndex) {
        return capture(fileUrl, new int[]{frameIndex})[0];
    }

    /**
     * 批量获取指定索引下的视频截图
     *
     * @param fileUrl      视频文件地址 本地文件地址，或者是网络地址
     * @param frameIndexes 要获取的帧序号数组 帧序号>=0的值 如果指定的序号位置没有图片，将使用后面最近的一个。
     * @return 帧图片对象数组
     */
    public static VideoCaptureImageResult[] capture(String fileUrl, int[] frameIndexes) {
        if (null == frameIndexes || frameIndexes.length == 0) {
            return null;
        }
        FFmpegFrameGrabber ffg = null;
        try {
            ffg = new FFmpegFrameGrabber(fileUrl);
            ffg.start();

            VideoCaptureImageResult[] images = new VideoCaptureImageResult[frameIndexes.length];
            for (int i = 0; i < frameIndexes.length; i++) {
                images[i] = capture(ffg, frameIndexes[i]);
            }
            return images;
        } catch (FrameGrabber.Exception e) {
            LOGGER.error("获取视频对象异常", e);
            return null;
        } finally {
            if (null != ffg) {
                try {
                    ffg.stop();
                } catch (FrameGrabber.Exception e) {
                    LOGGER.warn("关闭FFmpegFrameGrabber异常");
                }
            }
        }
    }

    /**
     * 批量获取视频截图
     *
     * @param fileUrl        视频文件地址 本地文件地址，或者是网络地址
     * @param fromFrameIndex 要获取的帧序号 >=0的值 如果指定的序号位置没有图片，将使用后面最近的一个。
     * @param counts         要获取的略图数
     * @param step           获取下一个截图时路过的帧数>0的值 1表示继续下一个，2表示路过一个。
     * @return 帧图片对象数组
     */
    public static VideoCaptureImageResult[] capture(String fileUrl, int fromFrameIndex, int counts, int step) {
        if (fromFrameIndex < 0) {
            fromFrameIndex = 0;
        }
        if (step < 1) {
            step = 1;
        }
        if (counts < 1) {
            counts = 1;
        }

        FFmpegFrameGrabber ffg = null;
        try {
            ffg = new FFmpegFrameGrabber(fileUrl);
            ffg.start();

            VideoCaptureImageResult[] images = new VideoCaptureImageResult[counts];
            int idx = 0; // 当前帧位置
            int i = 0; // 已经获取的帧图数
            int skipped = step; //第一个要，所属默认skipped=step，取第一个成功匹配的图
            VideoCaptureImageResult result = null;

            int frameLengths = ffg.getLengthInFrames();
            LOGGER.debug("总的帧数：{}", frameLengths);
            Frame frame = null;
            do {
                frame = ffg.grabFrame();
                // 起始位置确定
                while (idx < fromFrameIndex) {
                    LOGGER.trace("--跳过视频帧序号{}位置，开始位置指定！", idx);
                    frame = ffg.grabFrame();
                    idx++;
                }
                // 图片位置确定
                while (null == frame.image && idx < frameLengths) {
                    LOGGER.warn("--在视频帧序号{}位置，没有图片！", idx);
                    frame = ffg.grabFrame();
                    idx++;
                }

                boolean hasNext = null != frame.image && idx < frameLengths;
                result = getVideoCaptureImageResult(idx, idx, frame, hasNext);

                if (skipped >= step) {
                    images[i] = result;
                    LOGGER.debug("++使用获取的第{}个帧图, 帧序号：{}", i, idx);
                    i++;
                    skipped = 0;
                } else {
                    LOGGER.debug("--跳过视频帧序号{}位置，步长指定", idx);
                    skipped++;
                }
                idx++;

            } while (idx < frameLengths && i < counts && result.isHasNext());
            return images;
        } catch (FrameGrabber.Exception e) {
            LOGGER.error("获取视频对象异常", e);
            return null;
        } finally {
            if (null != ffg) {
                try {
                    ffg.stop();
                } catch (FrameGrabber.Exception e) {
                    LOGGER.warn("关闭FFmpegFrameGrabber异常");
                }
            }
        }
    }

    /**
     * 截取视频图片
     * <p>
     * 使用此方法需要确保ffmGrabber已经调用过start方法。
     *
     * @param ffmGrabber FFmpeg视频帧获取工具类对象
     * @param frameIndex 要获取的帧序号 >=0的值 如果指定的序号位置没有图片，将使用后面最近的一个。
     * @return 帧图片对象结果 未能获取到时返回null
     */
    public static VideoCaptureImageResult capture(FFmpegFrameGrabber ffmGrabber, int frameIndex) {
        try {
            if (frameIndex < 0) {
                frameIndex = 0;
            }
            int frameLengths = ffmGrabber.getLengthInFrames();

            if (frameIndex >= frameLengths) {
                LOGGER.error("指定的帧序号超出最大长度，不能获取！");
                return null;
            }

            int idx = 0;
            Frame frame = null;
            boolean breakFlag = false;
            while (idx < frameLengths) {
                frame = ffmGrabber.grabFrame();
                if (idx >= frameIndex
                        && null != frame.image) {
                    breakFlag = true;
                    break;
                }
                idx++;
            }
            LOGGER.debug("实际使用的帧序号：{}， 期望序号：{}", idx, frameIndex);
            if (null == frame) {
                LOGGER.error("未获取到帧数据");
                return null;
            }

            boolean hasNext = breakFlag || idx < frameIndex;

            return getVideoCaptureImageResult(frameIndex, idx, frame, hasNext);
        } catch (FrameGrabber.Exception e) {
            LOGGER.error("获取视频帧异常", e);
        }

        return null;
    }

    private static VideoCaptureImageResult getVideoCaptureImageResult(int needIdx, int idx, Frame frame, boolean hasNext) {
        BufferedImage image = null;
        if (null != frame.image) {
            Java2DFrameConverter java2DFrameConverter = new Java2DFrameConverter();
            image = java2DFrameConverter.convert(frame);
        }
        return new VideoCaptureImageResult(image, needIdx, idx, hasNext);
    }
}
